clear all; close all; clc;

%% Parameters
nu = 1;              % Diffusion coefficient
L = 1;               % Domain length
h = 1/20;            % Spatial step size
x = 0:h:L;           % Spatial grid
Nx = length(x);    % Number of interior points

% Time parameters
r_values = [1, 1/(2*h)];   % r = nu*dt/h^2 values to test
k_values = r_values*h^2/nu; % Corresponding time steps

% Number of steps to plot
steps_to_plot = [1, 2, 10];

%% Initial condition (triangular function)
phi = @(x) (x>=9/20 & x<1/2).*(20*(x-9/20)) + ...
           (x>=1/2 & x<11/20).*(-20*(x-11/20));
u0 = phi(x(1:end))'; % Interior points only

%% Compute and plot solutions
figure('Position', [100, 100, 1200, 600]);
titles = {'After 1 step', 'After 2 steps', 'After 10 steps'};
methods = {'G-L RK method r=1', 'G-L RK method r=1/2h'};

for col = 1:3
    steps = steps_to_plot(col);
    
    % Gauss-Legendre RK method r=1 (first row)
    subplot(2, 3, col);
    r = 1;
    k = r*h^2/nu;
    u = u0;
    for n = 1:steps
        u = GL(u, r, Nx);
    end
    plot(x(1:end), u, 'b-', 'LineWidth', 2);
    title(titles{col});
    ylim([0.0 0.6]);
    if col == 1
        ylabel(methods{1});
    end
    grid on;
    
    
    % Gauss-Legendre RK method r=1/2h (second row)
    subplot(2, 3, col+3);
    r = 1/(2*h);
    k = r*h^2/nu;
    u = u0;
    for n = 1:steps
        u = GL(u, r, Nx);
    end
    plot(x(1:end), u, 'g-', 'LineWidth', 2);
    ylim([-0.6 0.8]);
    if col == 1
        ylabel(methods{2});
    end
    grid on;
end

sgtitle('G-L RK method for Heat Equation');

function u_new = GL(u_old, r, Nx)
    % 1-stage Gauss-Legendre RK method for heat equation
    % Butcher tableau:
    % 1/2 | 1/2
    % --------
    %     | 1
    
    h = 1/20; 
    k = r*h^2; % Time step size
    
    % Create spatial difference operator matrix (centered difference)
    main_diag = -2*ones(Nx,1);
    off_diag = 1*ones(Nx-1,1);
    Dxx = (1/h^2) * (diag(main_diag) + diag(off_diag,1) + diag(off_diag,-1));
    
    % Implicit stage calculation (solve for K1)
    I = eye(Nx);
    K1 = (I - 0.5*k*Dxx) \ (Dxx*u_old);
    
    % Update solution
    u_new = u_old + k*K1;
end